# 基于Ma_q + b = \tau _q，给定tau_q, M, b，求解aq
import pinocchio as pin
from pinocchio.utils import rand, zero
import meshcat

# 引入src中所有的函数
import sys
from os.path import dirname, join

sys.path.append(join(dirname(dirname(__file__)), "src"))
sys.path.append(dirname(dirname(__file__)))
from src import *
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 获取URDF模型路径
URDF_FILE_PATH = filePath(URDF_NAME = 'ur5_gripper', ROBOT = True)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 构建robot对象，并得到viz
robot, viz = robot_viz(URDF_FILE_PATH)
robot.model.gravity.linear[:2] = 0.0
robot.model.gravity.linear[2] = -9.8
q = pin.randomConfiguration(robot.model)
vq = rand(robot.model.nv)
aq = zero(robot.model.nv)
t_q = zero(robot.model.nv)
pin.framesForwardKinematics(robot.model, robot.data, q)
pin.updateFramePlacements(robot.model, robot.data) 
robot.initViewer(loadModel=True)
viz.display(q)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

import math
import time

dt = 0.001

starttime = time.time()
for i in range(10000):

    # 设置阻尼力矩
    t_q = -0.5*vq
    # t_q += 10*math.sin(time.time()-starttime)


    # 机器人动力学方程为Ma_q + b = \tau _q, 此处b包含了Coriolis, centrifugal, gravity
    # rnea是recursive Newton Euler Algorithm的简写，即递归牛顿欧拉算法
    b = pin.rnea(robot.model, robot.data, q, vq, aq)


    # crba是专门计算机器人质量矩阵M的函数
    M = pin.crba(robot.model, robot.data, q)

    # 求aq， 不要用np.linalg.solve， 经常会不收敛
    aq = pin.aba(robot.model, robot.data, q, vq, t_q)
    vq += aq*dt
    # 求解新位置
    q = pin.integrate(robot.model, q, vq*dt)
    # 更新位置，这个一定不能忘
    pin.forwardKinematics(robot.model, robot.data, q)
    
    # 更新显示
    viz.display(q)